import torch
import numpy as np
from models import NightPedestrianPredictor
import yaml
import time

class PedestrianPredictor:
    def __init__(self, model_path, config_path):
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
        
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = NightPedestrianPredictor(self.config).to(self.device)
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        self.model.eval()
        
        self.history_buffer = []
        self.max_history = self.config['training']['history_frames']
    
    def process_frame(self, infrared_frame, radar_points, context):
        """处理单帧数据"""
        # 添加到历史缓冲区
        self.history_buffer.append({
            'infrared': infrared_frame,
            'radar': radar_points,
            'context': context
        })
        
        if len(self.history_buffer) > self.max_history:
            self.history_buffer.pop(0)
        
        # 如果缓冲区不足，返回空结果
        if len(self.history_buffer) < self.min_history_frames:
            return None, None
        
        # 准备模型输入
        batch_data = self.prepare_batch()
        
        with torch.no_grad():
            intent_probs, trajectories, tracks = self.model(*batch_data)
        
        return {
            'intent_probabilities': intent_probs.cpu().numpy(),
            'future_trajectories': trajectories.cpu().numpy(),
            'current_tracks': tracks.cpu().numpy()
        }
    
    def prepare_batch(self):
        """从历史缓冲区准备批次数据"""
        infrared_frames = []
        radar_frames = []
        context_frames = []
        
        for frame_data in self.history_buffer:
            infrared_frames.append(frame_data['infrared'])
            radar_frames.append(frame_data['radar'])
            context_frames.append(frame_data['context'])
        
        # 转换为tensor
        infrared_batch = torch.FloatTensor(np.array(infrared_frames)).unsqueeze(0).to(self.device)
        radar_batch = torch.FloatTensor(np.array(radar_frames)).unsqueeze(0).to(self.device)
        
        return infrared_batch, radar_batch, context_frames[0]  # 使用最新上下文

# 使用示例
if __name__ == '__main__':
    predictor = PedestrianPredictor('best_model.pth', 'configs/model_config.yaml')
    
    # 模拟输入数据
    infrared_frame = np.random.randn(1, 512, 640)  # 模拟红外帧
    radar_points = np.random.randn(100, 4)  # 模拟雷达点云
    context = {
        'light_intensity': 0.1,
        'signal_state': 1.0,
        'in_crosswalk': 0.0
    }
    
    result = predictor.process_frame(infrared_frame, radar_points, context)
    print("预测结果:", result)